package com.starhub.integration.sdk.shangtang.model;

import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.databind.ObjectMapper;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.model.StreamingResponseHandler;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.output.FinishReason;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.model.output.TokenUsage;
import okhttp3.*;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManager;
import javax.net.ssl.X509TrustManager;
import java.security.SecureRandom;
import java.security.cert.X509Certificate;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicReference;

import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.model.StreamingResponseHandler;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.output.FinishReason;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.model.output.TokenUsage;
import okhttp3.*;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManager;
import javax.net.ssl.X509TrustManager;
import java.security.SecureRandom;
import java.security.cert.X509Certificate;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

import static org.eclipse.jetty.http.HttpTester.parseResponse;

public class ShangtangStreamingChatModel implements StreamingChatLanguageModel {

    private static final Logger log = LoggerFactory.getLogger(ShangtangStreamingChatModel.class);

    private final String apiUrl;
    private final OkHttpClient httpClient;
    private final ObjectMapper objectMapper;
    private final String modelName;
    private final double temperature;
    private final int maxTokens;
    private final ExecutorService streamingExecutor;

    public ShangtangStreamingChatModel(String apiUrl, String modelName, double temperature, int maxTokens) {
        this.apiUrl = apiUrl;
        this.modelName = modelName;
        this.temperature = temperature;
        this.maxTokens = maxTokens;

        // 配置ObjectMapper忽略未知字段
        this.objectMapper = new ObjectMapper();
        this.objectMapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);

        // 创建绕过SSL验证的HTTP客户端
        this.httpClient = createUnsafeOkHttpClient();

        // 创建流式处理线程池
        this.streamingExecutor = Executors.newCachedThreadPool();
    }

    // 创建绕过SSL验证的HTTP客户端
    private OkHttpClient createUnsafeOkHttpClient() {
        try {
            // 1. 创建信任所有证书的TrustManager
            TrustManager[] trustAllCerts = new TrustManager[]{
                    new X509TrustManager() {
                        public void checkClientTrusted(X509Certificate[] chain, String authType) {}
                        public void checkServerTrusted(X509Certificate[] chain, String authType) {}
                        public X509Certificate[] getAcceptedIssuers() {
                            return new X509Certificate[]{};
                        }
                    }
            };

            // 2. 创建SSLContext
            SSLContext sslContext = SSLContext.getInstance("SSL");
            sslContext.init(null, trustAllCerts, new SecureRandom());

            // 3. 创建OkHttpClient
            return new OkHttpClient.Builder()
                    .sslSocketFactory(sslContext.getSocketFactory(), (X509TrustManager) trustAllCerts[0])
                    .hostnameVerifier((hostname, session) -> true) // 忽略主机名验证
                    .connectTimeout(120, TimeUnit.SECONDS)
                    .readTimeout(120, TimeUnit.SECONDS)
                    .build();
        } catch (Exception e) {
            throw new RuntimeException("创建HTTP客户端失败", e);
        }
    }

    @Override
    public void generate(List<ChatMessage> messages, StreamingResponseHandler<AiMessage> handler) {
        streamingExecutor.submit(() -> {
            try {
                ShangtangRequest request = new ShangtangRequest(
                        modelName,
                        convertMessages(messages),
                        true,
                        temperature,
                        maxTokens
                );

                log.info("商汤请求: {}", objectMapper.writeValueAsString(request));
                Request httpRequest = new Request.Builder()
                        .url(apiUrl)
                        .post(RequestBody.create(
                                objectMapper.writeValueAsString(request),
                                MediaType.get("application/json")
                        ))
                        .build();

                log.info("调用商汤API: {}", apiUrl);

                httpClient.newCall(httpRequest).enqueue(new Callback() {
                    @Override
                    public void onResponse(@NotNull Call call, @NotNull okhttp3.Response response) throws IOException {
                        if (!response.isSuccessful()) {
                            String errorBody = response.body() != null ? response.body().string() : "无响应体";
                            log.error("商汤API请求失败: {} - {}", response.code(), errorBody);
                            handler.onError(new RuntimeException("商汤API请求失败: " + response.code()));
                            return;
                        }
                        processStreamingResponse(response, handler);
                    }

                    @Override
                    public void onFailure(Call call, IOException e) {
                        log.error("商汤API调用失败", e);
                        handler.onError(e);
                    }
                });
            } catch (Exception e) {
                log.error("生成请求失败", e);
                handler.onError(e);
            }
        });
    }

    private void processStreamingResponse(okhttp3.Response response, StreamingResponseHandler<AiMessage> handler) {
        try (ResponseBody body = response.body();
             BufferedReader reader = new BufferedReader(new InputStreamReader(body.byteStream()))) {

            StringBuilder fullContent = new StringBuilder();
            TokenUsage tokenUsage = null;
            boolean receivedCompletionEvent = false;

            String line;
            while ((line = reader.readLine()) != null) {
                if (line.startsWith("event: ")) {
                    String eventType = line.substring(7).trim();
                    if ("error".equals(eventType)) {
                        handleErrorEvent(reader, handler);
                        return;
                    }
                } else if (line.startsWith("data: ")) {
                    String json = line.substring(6).trim();

                    if ("[DONE]".equals(json)) {
                        receivedCompletionEvent = true;
                        break;
                    }

                    try {
                        ShangtangResponse shangtangResponse = parseResponse(json);

                        if (isFullContentEvent(shangtangResponse)) {
                            log.debug("忽略完整内容事件");
                            continue;
                        }

                        if (shangtangResponse.getChoices() != null && !shangtangResponse.getChoices().isEmpty()) {
                            ShangtangResponse.Choice choice = shangtangResponse.getChoices().get(0);
                            if (choice.getDelta() != null &&
                                    choice.getDelta().getContent() != null &&
                                    !choice.getDelta().getContent().isEmpty()) {

                                String token = choice.getDelta().getContent();
                                handler.onNext(token);
                                fullContent.append(token);
                            }
                        }

                        if (shangtangResponse.getUsage() != null) {
                            tokenUsage = new TokenUsage(
                                    shangtangResponse.getUsage().getPrompt_tokens(),
                                    shangtangResponse.getUsage().getCompletion_tokens(),
                                    shangtangResponse.getUsage().getTotal_tokens()
                            );
                        }
                    } catch (Exception e) {
                        log.warn("解析商汤响应失败: {} - {}", json, e.getMessage());
                    }
                }
            }

            // 修改完成处理逻辑：即使未收到 [DONE] 事件也视为完成
            if (fullContent.length() > 0) {
                handler.onComplete(new Response<>(
                        AiMessage.from(fullContent.toString()),
                        tokenUsage,
                        FinishReason.STOP
                ));
                log.debug("流式响应正常结束（未收到DONE事件）");
            } else if (receivedCompletionEvent) {
                handler.onComplete(new Response<>(
                        AiMessage.from(fullContent.toString()),
                        tokenUsage,
                        FinishReason.STOP
                ));
            } else {
                log.error("流式响应未正常结束且无有效内容");
                handler.onError(new RuntimeException("流式响应未正常结束"));
            }

        } catch (Exception e) {
            log.error("处理流式响应失败", e);
            handler.onError(e);
        }
    }
    // 判断是否是完整内容事件
    private boolean isFullContentEvent(ShangtangResponse response) {
        if (response.getChoices() == null || response.getChoices().isEmpty())
            return false;

        ShangtangResponse.Delta delta = response.getChoices().get(0).getDelta();
        return delta != null &&
                delta.getContent() != null &&
                delta.getContent().length() > 200; // 阈值根据实际情况调整
    }

    // 处理错误事件
    private void handleErrorEvent(BufferedReader reader, StreamingResponseHandler<AiMessage> handler) {
        try {
            String errorData = "";
            String line;
            while ((line = reader.readLine()) != null) {
                if (line.startsWith("data: ")) {
                    errorData = line.substring(6).trim();
                    break;
                }
            }
            handler.onError(new RuntimeException("商汤API错误: " + errorData));
        } catch (IOException ex) {
            handler.onError(new RuntimeException("读取错误信息失败", ex));
        }
    }

    // 转换消息格式
    private List<ShangtangMessage> convertMessages(List<ChatMessage> messages) {
        return messages.stream()
                .map(msg -> new ShangtangMessage(
                        msg.type().name().toLowerCase(),
                        msg.text()
                ))
                .collect(java.util.stream.Collectors.toList());
    }

    // 解析响应
    private ShangtangResponse parseResponse(String json) throws IOException {
        return objectMapper.readValue(json, ShangtangResponse.class);
    }

    // 关闭资源
    public void close() {
        streamingExecutor.shutdown();
        try {
            if (!streamingExecutor.awaitTermination(5, TimeUnit.SECONDS)) {
                streamingExecutor.shutdownNow();
            }
        } catch (InterruptedException e) {
            streamingExecutor.shutdownNow();
            Thread.currentThread().interrupt();
        }
    }

    // 商汤API请求结构
    private static class ShangtangRequest {
        private final String model;
        private final List<ShangtangMessage> messages;
        private final boolean stream;
        private final double temperature;
        private final int max_tokens;

        public ShangtangRequest(String model, List<ShangtangMessage> messages,
                                boolean stream, double temperature, int max_tokens) {
            this.model = model;
            this.messages = messages;
            this.stream = stream;
            this.temperature = temperature;
            this.max_tokens = max_tokens;
        }

        // Getters
        public String getModel() { return model; }
        public List<ShangtangMessage> getMessages() { return messages; }
        public boolean isStream() { return stream; }
        public double getTemperature() { return temperature; }
        public int getMax_tokens() { return max_tokens; }
    }

    // 消息结构
    private static class ShangtangMessage {
        private final String role;
        private final String content;

        public String getRole() {
            return role;
        }

        public String getContent() {
            return content;
        }

        public ShangtangMessage(String role, String content) {
            this.role = role;
            this.content = content;
        }
    }

    // 商汤API响应结构
    @JsonIgnoreProperties(ignoreUnknown = true)
// 商汤API响应结构
    private static class ShangtangResponse {
        private String id;
        private String object;
        private long created;
        private String model;
        private List<Choice> choices;
        private Usage usage;

        // Getters and setters
        public String getId() { return id; }
        public void setId(String id) { this.id = id; }

        public String getObject() { return object; }
        public void setObject(String object) { this.object = object; }

        public long getCreated() { return created; }
        public void setCreated(long created) { this.created = created; }

        public String getModel() { return model; }
        public void setModel(String model) { this.model = model; }

        public List<Choice> getChoices() { return choices; }
        public void setChoices(List<Choice> choices) { this.choices = choices; }

        public Usage getUsage() { return usage; }
        public void setUsage(Usage usage) { this.usage = usage; }

        public static class Choice {
            private String finish_reason;
            private int index;
            private Delta delta;

            // Getters and setters
            public String getFinish_reason() { return finish_reason; }
            public void setFinish_reason(String finish_reason) { this.finish_reason = finish_reason; }

            public int getIndex() { return index; }
            public void setIndex(int index) { this.index = index; }

            public Delta getDelta() { return delta; }
            public void setDelta(Delta delta) { this.delta = delta; }
        }

        public static class Delta {
            private String content;
            private String role;

            // Getters and setters
            public String getContent() { return content; }
            public void setContent(String content) { this.content = content; }

            public String getRole() { return role; }
            public void setRole(String role) { this.role = role; }
        }

        public static class Usage {
            private int prompt_tokens;
            private int completion_tokens;
            private int total_tokens;

            // Getters and setters
            public int getPrompt_tokens() { return prompt_tokens; }
            public void setPrompt_tokens(int prompt_tokens) { this.prompt_tokens = prompt_tokens; }

            public int getCompletion_tokens() { return completion_tokens; }
            public void setCompletion_tokens(int completion_tokens) { this.completion_tokens = completion_tokens; }

            public int getTotal_tokens() { return total_tokens; }
            public void setTotal_tokens(int total_tokens) { this.total_tokens = total_tokens; }
        }
    }
}
